from datetime import datetime, time as datetime_time, timedelta
import time
import sys
import os

# Define the data structure for the train information
# Note: the departure and arrival stations are station names
class TrainInfo_s:
	def __init__(self, train_no, dept_city, arrv_city, dept_station, arrv_station, dept_time, arrv_time, duration, op_days):
		self.train_no = train_no			# Type: string
		self.dept_city = dept_city			# Type: string (of city name)
		self.arrv_city = arrv_city
		self.dept_station = dept_station	# Type: string (of station name)
		self.arrv_station = arrv_station
		self.dept_time = dept_time			# Type: datetime_time
		self.arrv_time = arrv_time			# Type: datetime_time
		self.duration = duration			# Type: datetime_time
		self.op_days = op_days				# Type: list[int]
	def is_power_distributed_EMU(self):
		ret = False
		if self.train_no[0] == 'G':
			ret = True
		elif self.train_no[0] == 'D':
			# Convert the train number to an integer
			train_no_int = int(self.train_no[1:])
			if train_no_int > 300:
				ret = True
			# Otherwise, it is a power-concentrated EMU (CR200J)
		elif self.train_no[0] == 'C':
			train_no_int = int(self.train_no[1:])
			if train_no_int > 1000:
				ret = True
			# Otherwise, it is a power-concentrated EMU
		return ret

# Reduced train information class
class TrainInfo_r:
	def __init__(self, train_no, dept_city, arrv_city, duration, op_days):
		self.train_no = train_no		# Type: string
		self.dept_city = dept_city		# Type: string (of city name)
		self.arrv_city = arrv_city
		self.duration = duration		# Type: `datetime_time`
		self.op_days = op_days			# Type: integer

# Read the detailed CSV file and return a list of `TrainInfo_s` objects
# Note: all of the trains have the same departure cities
def read_detailed_csv(file_name):
	train_list = []
	# Use UTF-8 encoding to read the file
	with open(file_name, "r", encoding="utf-8") as file:
		lines = file.readlines()
		# Skip the first line
		lines = lines[1:]
		for line in lines:
			train_info = line.split(",")
			train_no = train_info[0]
			dept_city = train_info[1]
			arrv_city = train_info[2]
			dept_station = train_info[3]
			arrv_station = train_info[4]
			dept_time = datetime.strptime(train_info[5], "%H:%M").time()
			arrv_time = datetime.strptime(train_info[6], "%H:%M").time()
			duration = datetime.strptime(train_info[7], "%H:%M").time()
			# `train_info[8:14]` is a list of '1' or '0'
			op_days = [int(x) for x in train_info[8:15]]
			train_list.append(TrainInfo_s(train_no, dept_city, arrv_city, dept_station, arrv_station, dept_time, arrv_time, duration, op_days))
	return train_list

# Read the detailed train information (`train_list` is a list of `TrainInfo_s` that has the same departure city) and return a list of `TrainInfo_r` objects
def reduce_train_info(train_list):
	train_list_reduced = []
	for train in train_list:
		# If the train is NOT a power-distributed EMU, skip it
		if not train.is_power_distributed_EMU():
			continue
		# Search `train_list_reduced` for the same train number and arrival city
		flag_found = False
		for t in train_list_reduced:
			if t.train_no == train.train_no and t.arrv_city == train.arrv_city:
				flag_found = True
				# If found and the duration is shorter, update the duration and operating days
				if t.duration > train.duration:
					t.duration = train.duration
					t.op_days = sum(train.op_days)
				break
		if not flag_found:
			# If not found, append a new train
			train_list_reduced.append(TrainInfo_r(train.train_no, train.dept_city, train.arrv_city, train.duration, sum(train.op_days)))
	return train_list_reduced

def main():
	# Read the command line arguments
	if len(sys.argv) != 2:
		print("Usage: python DataReduction.py <path_of_data>")
		exit(1)
	data_path = sys.argv[1]
	# Read every CSV file in the directory
	for file_name in os.listdir(data_path):
		if file_name.endswith("-d.csv"):
			file_name = os.path.join(data_path, file_name)
			train_list = read_detailed_csv(file_name)
			train_list_reduced = reduce_train_info(train_list)
			# Write the reduced train information to a new CSV file
			with open(file_name.replace("-d.csv", "-r.csv"), "w", encoding="utf-8") as file:
				file.write("train_no,dept_city,arrv_city,duration,op_days\n")
				for train in train_list_reduced:
					# Convert the duration to a string
					duration_str = train.duration.strftime("%H:%M")
					file.write(f"{train.train_no},{train.dept_city},{train.arrv_city},{duration_str},{train.op_days}\n")

if __name__ == '__main__':
	main()
